#!/usr/bin/env python
"""
MNIST训练脚本 - 支持重复实验和多尺度分支控制版本（中文日志）
"""

import os
import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm
from pathlib import Path
import sys
import argparse

# --- 1. 路径初始化 ---
sys.path.append(str(Path(__file__).parent.parent))

from core.environment import Environment
from data.mnist_dataset import MNISTDataModule
from models.factory import create_model
from utils.logger import setup_logger
from utils.config_loader import load_config
from trainers.base_trainer import BaseTrainer

# --- 扩展命令行参数：增加实验次数和分支控制 ---
parser = argparse.ArgumentParser(description='MNIST模型对比实验（支持重复实验和分支控制）')
parser.add_argument('--model', type=str, required=True,
                    choices=['simple_cnn', 'lenet5', 'multi_scale_cnn'],
                    help='指定模型名称')
parser.add_argument('--runs', type=int, default=1,
                    help='重复实验轮次（默认1次）')

# 多尺度CNN专用参数（控制分支启用）
parser.add_argument("--use_branch1", type=lambda x: x.lower() == "true", default=True,
                    help="是否使用3×3分支 (true/false，默认true)")
parser.add_argument("--use_branch2", type=lambda x: x.lower() == "true", default=True,
                    help="是否使用5×5分支 (true/false，默认true)")
parser.add_argument("--use_branch3", type=lambda x: x.lower() == "true", default=True,
                    help="是否使用7×7分支 (true/false，默认true)")

args = parser.parse_args()


def main():
    # 全局配置加载
    env_type = os.getenv('RUN_ENV', 'local')
    config = load_config(env_type)
    env = Environment(config)

    # --- 2. 准备多轮实验的基础路径 ---
    # 多尺度模型：在目录名中体现分支配置，便于区分不同消融实验
    if args.model == "multi_scale_cnn":
        branch_suffix = f"_b1-{args.use_branch1}_b2-{args.use_branch2}_b3-{args.use_branch3}"
        model_root = env.project_root / f"results_{args.model}{branch_suffix}"
    else:
        model_root = env.project_root / f"results_{args.model}"

    model_root.mkdir(exist_ok=True)

    # --- 3. 循环执行多轮实验 ---
    for run_id in range(args.runs):
        print(f"\n===== 开始实验轮次 {run_id + 1}/{args.runs} =====")

        # 为每轮实验创建独立目录
        run_dir = model_root / f"run_{run_id}"
        run_log_dir = run_dir / "日志"
        run_ckpt_dir = run_dir / "模型权重"
        run_log_dir.mkdir(parents=True, exist_ok=True)
        run_ckpt_dir.mkdir(parents=True, exist_ok=True)

        # 初始化本轮实验的日志
        logger = setup_logger(
            name=f"mnist_train_{args.model}_run{run_id}",
            log_dir=run_log_dir
        )
        logger.info(f"实验轮次 {run_id + 1} 初始化完成")
        logger.info(f"结果保存路径: {run_dir}")
        # 记录分支配置（多尺度模型专用）
        if args.model == "multi_scale_cnn":
            logger.info(f"分支配置: b1={args.use_branch1}, b2={args.use_branch2}, b3={args.use_branch3}")

        # --- 4. 数据准备 ---
        datamodule = MNISTDataModule(
            batch_size=env.batch_size,
            data_dir=env.project_root / "data"
        )
        datamodule.prepare_data()
        datamodule.setup()
        dataloaders = datamodule.get_dataloaders()
        train_loader = dataloaders['train']
        val_loader = dataloaders['val']

        # --- 5. 模型与优化器（每轮重新初始化）---
        # 多尺度模型：传入分支控制参数
        if args.model == "multi_scale_cnn":
            model = create_model(
                args.model,
                num_classes=10,
                use_branch1=args.use_branch1,
                use_branch2=args.use_branch2,
                use_branch3=args.use_branch3
            ).to(env.device)
        else:
            model = create_model(args.model, num_classes=10).to(env.device)

        optimizer = optim.Adam(
            model.parameters(),
            lr=env.lr,
            weight_decay=float(env.weight_decay)
        )
        scheduler = optim.lr_scheduler.StepLR(
            optimizer,
            step_size=5,
            gamma=0.1
        )

        # 初始化训练器
        trainer = BaseTrainer(
            model=model,
            optimizer=optimizer,
            scheduler=scheduler,
            device=env.device,
            patience=5
        )

        # --- 6. 训练循环 ---
        best_val_acc = 0.0  # 记录本轮最佳准确率
        for epoch in range(1, env.epochs + 1):
            # 训练阶段
            train_loss = trainer.train_epoch(train_loader)
            # 验证阶段
            val_acc, early_stop = trainer.validate(val_loader)
            # 日志记录（使用中文）
            logger.log_metrics({
                '训练损失': train_loss,
                '验证准确率': val_acc
            }, epoch)
            # 保存本轮最佳模型
            if val_acc > best_val_acc:
                best_val_acc = val_acc
                torch.save(
                    model.state_dict(),
                    run_ckpt_dir / f"最佳模型_epoch{epoch}.pth"
                )
                logger.info(f"轮次 {run_id + 1} 最佳模型更新（准确率: {best_val_acc:.4f}）")
            # 早停检查
            if early_stop:
                logger.info(f"轮次 {run_id + 1} 早停于第 {epoch} 轮")
                break

        logger.info(f"轮次 {run_id + 1} 完成，最佳验证准确率: {best_val_acc:.4f}")

    print(f"\n所有 {args.runs} 轮实验已完成，结果保存于: {model_root}")


if __name__ == "__main__":
    main()
